#!/bin/bash

export ARCH_AND_GPU_INFO="cpu"

install_torch_win() {
    PYTORCH_CUDA_PREFIX=$1
    export TMP_TORCH_VERSION=""
    local TORCH_INDEX_URL="https://download.pytorch.org/whl/nightly/cpu"
    if ! command -v nvidia-smi &> /dev/null; then
        echo "No NVIDIA driver detected."
        if command -v rocm-smi &> /dev/null && rocm-smi | grep -i "amd" > /dev/null; then
            echo "Found AMD GPU"
            TORCH_INDEX_URL=https://download.pytorch.org/whl/rocm6.2
        else
            echo "only cpu"
            TORCH_INDEX_URL=""
        fi
    else
        if ! command -v nvcc &> /dev/null; then
            echo "Error: The nvcc command was not found. Make sure you have the CUDA Toolkit installed. use index /cu128"
            TORCH_INDEX_URL="${PYTORCH_CUDA_PREFIX}/cu128"
            export TMP_TORCH_VERSION="==2.7.1"
        else
            echo "Found NVIDIA drivers and cuda"
            nvcc_output=$(nvcc --version)
            if [[ $? -ne 0 ]]; then
                echo "Error: Failed to execute 'nvcc --version'." >&2
                # exit 1
            fi

            cuda_line=$(echo "$nvcc_output" | grep 'Cuda compilation tools')

            if [[ $cuda_line =~ release\ ([0-9]+\.[0-9]+) ]]; then
                cuda_version=${BASH_REMATCH[1]}
                echo "CUDA Version: $cuda_version"
            else
                echo "Error: Unable to parse CUDA version from output." >&2
                exit 1
            fi

            if [ -z "$cuda_version" ]; then
                echo "Error: Unable to extract CUDA version number from nvcc output."
                TORCH_INDEX_URL="${PYTORCH_CUDA_PREFIX}/cu128"
                export TMP_TORCH_VERSION="==2.7.1"
            else
                if [[ "$cuda_version" == "12.8" || "$cuda_version" == "12.9" || "$cuda_version" == "13.0" || "$cuda_version" =~ ^13\.[1-9] ]]; then
                    TORCH_INDEX_URL="${PYTORCH_CUDA_PREFIX}/cu128"
                    export TMP_TORCH_VERSION="==2.7.1"
                elif [[ "$cuda_version" == "12.6" || "$cuda_version" == "12.7" ]]; then
                    TORCH_INDEX_URL="${PYTORCH_CUDA_PREFIX}/cu126"
                    export TMP_TORCH_VERSION="==2.7.1"
                elif [[ "$cuda_version" == "12.4" || "$cuda_version" == "12.5" ]]; then
                    TORCH_INDEX_URL="${PYTORCH_CUDA_PREFIX}/cu124"
                elif [[ "$cuda_version" == "12.0" || "$cuda_version" == "12.1" || "$cuda_version" == "12.2" || "$cuda_version" == "12.3" ]]; then
                    TORCH_INDEX_URL="${PYTORCH_CUDA_PREFIX}/cu121"
                else
                    TORCH_INDEX_URL="${PYTORCH_CUDA_PREFIX}/cu128"
                fi
                export ARCH_AND_GPU_INFO="nvidia_gpu"
            fi
        fi
    fi
    

    if [ -z "$TORCH_INDEX_URL" ]; then
        INDEX_PARAM=""
        echo "Do not set index url."
    else
        INDEX_PARAM="--index-url $TORCH_INDEX_URL"
    fi

    echo "Start install torch from: $TORCH_INDEX_URL"
    echo "INDEX_PARAM: $INDEX_PARAM"
    conda_run_pip_install torch${TMP_TORCH_VERSION} torchvision torchaudio $INDEX_PARAM
}

